import dolphindb.orca as orca
import matplotlib.pyplot as plt
import warnings
warnings.filterwarnings("ignore")

orca.connect('115.239.209.122', 25668, 'admin', '123456')

def main():
    # TAQ quotes
    print("Info of TAQ quotes")
    quotes = orca.read_table('dfs://TAQ', 'quotes')
    print("The total records is：", len(quotes))

    # calculate the average bid-ask spread per minute of LEH on 2007.08.31
    print("calculate the average bid-ask spread per minute of LEH on 2007.08.31")
    tmp = quotes[
        quotes.date == '2007.08.31',
        quotes.symbol == 'LEH',
        quotes.time.between('09:30:00', '15:59:59'),
        quotes.ofr > quotes.bid,
        quotes.bid > 0,
        quotes.ofr / quotes.bid < 1.2
    ]
    tmp.set_index('time', inplace=True)
    avg_spread = ((tmp.ofr - tmp.bid) / (tmp.ofr + tmp.bid) * 2).groupby(tmp.index.hourminute, lazy=True).mean()
    avg_spread.plot()
    plt.show()

    # calculate the average bid-ask spread per minute of all stocks on 2007.09.24
    print("calculate the average bid-ask spread per minute of all stocks on 2007.09.24")
    tmp = quotes[
        quotes.date == '2007.09.25',
        quotes.time.between('09:30:00', '15:59:59'),
        quotes.ofr > quotes.bid,
        quotes.bid > 0,
        quotes.ofr / quotes.bid < 1.2
    ]
    tmp.set_index('time', inplace=True)
    avg_spread = ((tmp.ofr - tmp.bid) / (tmp.ofr + tmp.bid) * 2).groupby(tmp.index.hourminute, lazy=True).mean()
    avg_spread.plot()
    plt.show()

    # calculate the average bid-ask spread per minute of all stocks on 2007.09.25, 2007.09.26, 2007.09.27
    print("calculate the average bid-ask spread per minute of all stocks on 2007.09.25, 2007.09.26, 2007.09.27")
    tmp = quotes[
        quotes.date.between('2007.09.25', '2007.09.27'),
        quotes.time.between('09:30:00', '15:59:59'),
        quotes.ofr > quotes.bid,
        quotes.bid > 0,
        quotes.ofr / quotes.bid < 1.2
    ]
    tmp.set_index('time', inplace=True)
    avg_spread = ((tmp.ofr - tmp.bid) / (tmp.ofr + tmp.bid) * 2).groupby(tmp.index.hourminute, lazy=True).mean()
    avg_spread.plot()
    plt.show()

    # calculate the average quote size spread per minute of all stocks on 2007.09.25, 2007.09.26, 2007.09.27
    print("calculate the average quote size spread per minute of all stocks on 2007.09.25, 2007.09.26, 2007.09.27")
    tmp = quotes[
        quotes.date.between('2007.09.25', '2007.09.27'),
        quotes.time.between('09:30:00', '15:59:59'),
        quotes.ofr > quotes.bid,
        quotes.bid > 0,
        quotes.ofr / quotes.bid < 1.2
    ]
    tmp.set_index('time', inplace=True)
    quote_size = (tmp.bidsiz + tmp.ofrsiz).groupby(tmp.index.hourminute, lazy=True).mean()
    quote_size.plot()
    plt.show()

    # TAQ trades
    print("Info of TAQ trades")
    trades = orca.read_table('dfs://TAQ', 'trades')
    print("The total records is：", len(trades))

    # calculate the total volme spread per minute of all stocks on 2007.09.25, 2007.09.26, 2007.09.27
    print("calculate the total volme spread per minute of all stocks on 2007.09.25, 2007.09.26, 2007.09.27")
    tmp = trades[
        trades.DATE.between('2007.09.25', '2007.09.27'),
        trades.TIME.between('09:30:00', '15:59:59')
    ]
    tmp.set_index('TIME', inplace=True)
    volumn = tmp.groupby(tmp.index.hourminute, lazy=True)['SIZE'].sum()
    volumn.plot()
    plt.show()

    # EQY
    print("Info of EQY")
    trade = orca.read_table('dfs://EQY', 'trade')
    nbbo = orca.read_table('dfs://EQY', 'nbbo')

    # calculate the average bid-ask spread per minute of all stocks on 2016.10.24
    print("calculate the average bid-ask spread per minute of all stocks on 2016.10.24")
    tmp = nbbo[
        nbbo.Time.dt.time.between('09:30:00', '15:59:59'),
        nbbo.Offer_Price > nbbo.Bid_Price,
        nbbo.Bid_Price > 0,
        nbbo.Offer_Price / nbbo.Bid_Price < 1.2,
    ]
    tmp.set_index('Time', inplace=True)
    avg_spread = ((tmp.Offer_Price - tmp.Bid_Price) / (tmp.Offer_Price + tmp.Bid_Price) * 2).groupby(
        tmp.index.hourminute, lazy=True).mean()
    avg_spread.plot()
    plt.show()

    # calculate the average quote size per minute of all stocks on 2016.10.24
    print("calculate the average quote size per minute of all stocks on 2016.10.24")
    quote_size = (tmp.Offer_Size + tmp.Bid_Size).groupby(tmp.index.hourminute, lazy=True).mean()
    quote_size.plot()
    plt.show()

    # calculate the total volume per minute of all stocks on 2016.10.24
    print("calculate the total volume per minute of all stocks on 2016.10.24")
    tmp = trade[trade.Time.dt.time.between('09:30:00', '15:59:59')]
    tmp.set_index('Time', inplace=True)
    volume = tmp.groupby(tmp.index.hourminute, lazy=True)['Trade_Volume'].sum()
    volume.plot()
    plt.show()

if __name__ == '__main__':
    main()